import torch
import torch.nn as nn


class DepthwiseSeparableConv(nn.Module):
    """
    Depthwise Separable Convolution:
    - Depthwise convolution: Applies a single filter per input channel (depth).
    - Pointwise convolution: Combines the outputs from depthwise convolution (1x1 Conv).
    """

    def __init__(self, in_channels, out_channels, stride):
        super(DepthwiseSeparableConv, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1,
                                   groups=in_channels, bias=False)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn_depthwise = nn.BatchNorm2d(in_channels)
        self.bn_pointwise = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.bn_depthwise(x)
        x = self.relu(x)
        x = self.pointwise(x)
        x = self.bn_pointwise(x)
        x = self.relu(x)
        return x


class MobileNetV1(nn.Module):
    def __init__(self, num_classes=1000):
        super(MobileNetV1, self).__init__()

        def conv_bn(in_channels, out_channels, stride):
            return nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True)
            )

        self.model = nn.Sequential(
            conv_bn(3, 32, stride=2),  # Conv1
            DepthwiseSeparableConv(32, 64, stride=1),  # Conv2_1
            DepthwiseSeparableConv(64, 128, stride=2),  # Conv2_2
            DepthwiseSeparableConv(128, 128, stride=1),  # Conv3_1
            DepthwiseSeparableConv(128, 256, stride=2),  # Conv3_2
            DepthwiseSeparableConv(256, 256, stride=1),  # Conv4_1
            DepthwiseSeparableConv(256, 512, stride=2),  # Conv4_2
            # Conv5_1 ~ Conv5_5 (5 blocks with stride=1)
            *[DepthwiseSeparableConv(512, 512, stride=1) for _ in range(5)],
            DepthwiseSeparableConv(512, 1024, stride=2),  # Conv5_6
            DepthwiseSeparableConv(1024, 1024, stride=1),  # Conv6
        )
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(1024, num_classes)

    def forward(self, x):
        x = self.model(x)
        x = self.avg_pool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x


# 测试网络
if __name__ == "__main__":
    model = MobileNetV1(num_classes=1000)

    print(f'Total parameters: {sum(param.numel() for param in model.parameters()) / 1e6:.2f} M')
    # 创建一个随机输入张量
    x = torch.randn(1, 3, 224, 224)  # Batch size=1, Channels=3, Height=224, Width=224
    output = model(x)
    print(output.shape)  # 应输出 torch.Size([1, 1000])
    # 特点：深度可分离卷积